% Updates G for the Multisector GE model
%
% Jon Steinsson and Emi Nakamura, July 2006
%****************************************************

function [G_new W_SPold supDifG G_temp GMove] = ...
    UpdateGMultiSectorCalvoPlus(Q,F,F2,G_old,alphaCalvo,a,rp,rad,theta,...
    Proba,ProbNgr,SectorWeights,GMove_old)

agridsize = size(a,1);
rpgridsize = size(rp,1);
radgridsize = size(rad,1);
Ssize = agridsize*rpgridsize*radgridsize;

NSectors = size(SectorWeights,1);

% Get weights as a function of S_t/P_t-1
%********************************************************

QaInd = repmat((1:agridsize)',rpgridsize*radgridsize,1);

QradInd = repmat((1:radgridsize)',[1 agridsize rpgridsize]);
QradInd = permute(QradInd,[2 3 1]);
QradInd = reshape(QradInd,[Ssize 1]);

OldInd = (1:Ssize)';

for ii = 1:NSectors
    [dummy,QrpInd] = ismember(int32(F(:,ii).F*10^6),int32(rp*10^6));
    clear dummy
    %QrpInd = match(int32(F(:,ii).F*10^6),int32(rp*10^6));
    FInd(ii).FInd = sub2ind([agridsize rpgridsize radgridsize],QaInd,QrpInd,QradInd);
    [FInd(ii).FInd iFInd] = sort(FInd(ii).FInd);
    OldFInd(ii).OldFInd = OldInd(iFInd);
    
    [dummy,QrpInd] = ismember(int32(F2(:,ii).F2*10^6),int32(rp*10^6));
    clear dummy
    %QrpInd = match(int32(F2(:,ii).F2*10^6),int32(rp*10^6));
    F2Ind(ii).F2Ind = sub2ind([agridsize rpgridsize radgridsize],QaInd,QrpInd,QradInd);
    [F2Ind(ii).F2Ind iF2Ind] = sort(F2Ind(ii).F2Ind);
    OldF2Ind(ii).OldF2Ind = OldInd(iF2Ind);
end
clear QaInd QrpInd1 QrpInd iFInd iF2ind OldInd

% Apply F
for ii = 1:NSectors
    Q_temp = accumarray(FInd(ii).FInd,Q(ii).Q(OldFInd(ii).OldFInd));
    Q_temp = [Q_temp; zeros(Ssize-size(Q_temp,1),1)];
    Q_temp2 = accumarray(F2Ind(ii).F2Ind,Q(ii).Q(OldF2Ind(ii).OldF2Ind));
    Q_temp2 = [Q_temp2; zeros(Ssize-size(Q_temp2,1),1)];
    Q(ii).Q = alphaCalvo(ii)*Q_temp + (1-alphaCalvo(ii))*Q_temp2;
end
clear Q_temp Q_temp2

% Apply Proba and ProbNgr
for ii = 1:NSectors
    Q(ii).Q = reshape(((reshape(Q(ii).Q,agridsize,rpgridsize*radgridsize)')*Proba(ii).Proba)',[Ssize 1]);
    Q(ii).Q = reshape(reshape(Q(ii).Q,[agridsize*rpgridsize radgridsize])*ProbNgr,[Ssize 1]);

end

% Get the prices
%**********************************************************

GInd = CalculateGInd(a,rp,rad,G_old);

for ii = 1:NSectors
    F(ii).F = F(ii).F(GInd);
end

% Create Price Index
%***************************************************************

totWeight = zeros(radgridsize,1);
sumF = zeros(radgridsize,1);
for ii = 1:NSectors
    F(ii).F = exp(F(ii).F).^(1-theta);
    Q(ii).Q = SectorWeights(ii)*Q(ii).Q;
    totWeight = totWeight ...
        + sum(reshape(Q(ii).Q,[agridsize*rpgridsize radgridsize]))';
    sumF = sumF + sum(reshape(F(ii).F.*Q(ii).Q,[agridsize*rpgridsize radgridsize]))';
end

% Deal with zero mass
W_zero = (totWeight < 10^(-10));
totWeight = totWeight + W_zero;
sumF_zero = (sumF < 10^(-10));
sumF = sumF + sumF_zero;

G_new = log((sumF./totWeight).^(1/(1-theta)));
G_new = (1-sumF_zero).*(1-W_zero).*G_new ...
    + W_zero.*zeros(radgridsize,1) ...
    + (1-W_zero).*sumF_zero.*zeros(radgridsize,1);
G_temp = G_new;
supDifG = max(abs(G_new));

rpinc = rp(2) - rp(1);
GstepH = zeros(radgridsize,1);
GstepH(2:end,1) = (abs(G_old(1:end-1,1)-G_old(2:end,1)) >  rpinc/2);
GstepH(1) = 1;

GstepL = zeros(radgridsize,1);
GstepL(1:end-1,1) = (abs(G_old(1:end-1,1)-G_old(2:end,1)) >  rpinc/2);
GstepL(end) = 1;

GmoveU = (G_new > rpinc/1.6).*GstepL;
GmoveD = (G_new < - rpinc/1.6).*GstepH;

GmoveUshift = [0; GmoveU(1:end-1,1)];
GmoveDshift = [GmoveD(2:end,1); 0];

% Create indicators for the two cases
MoveUp = (GmoveU & (1-GmoveDshift)); 
MoveDw = (GmoveD & (1-GmoveUshift));

% Get rid of cycling on many points
while sum(MoveUp) > radgridsize/5
    MoveUp = MoveUp.*(rand(radgridsize,1) > 0.25);
end
while sum(MoveDw) > radgridsize/5
    MoveDw = MoveDw.*(rand(radgridsize,1) > 0.25);
end

%[G_old G_new MoveUp MoveDw]

GMove = MoveUp*rpinc - MoveDw*rpinc;
% Get rid of cycling on few points
if max(abs(GMove+GMove_old)) < rpinc/2
    GMove = zeros(radgridsize,1);
end

G_new = G_old + GMove;

%supDifG = max(abs(G_new-G_old));

if max(abs(G_new-G_old)) < rpinc/2
    supDifG = 0;
end



% rpinc = rp(2) - rp(1);
% if StageForG == 1
%     %G_new = G_old + 1.1*rpinc*(G_new > rpinc) ...
%     %    - 1.1*rpinc*(G_new < -rpinc);
%     G_new = G_old + 1.1*rpinc*(G_new.*rand(radgridsize,1) > rpinc) ...
%         - 1.1*rpinc*(G_new.*rand(radgridsize,1) < -rpinc);
%     supDifG = max(abs(G_new-G_old));
%     if supDifG < rpinc
%         %supDifG = max(abs(G_new-G_old));
%         supDifG = 2*rpinc;
%         StageForG = 2;
%     end
% else
%     %[(G_new > rpinc) (G_new < -rpinc) (abs(G_new) == max(abs(G_new))*ones(size(G_new,1),1)) G_new]
%     G_new = G_old + 1.02*min(rpinc,max(abs(G_new)))*(G_new > 0).*(abs(G_new) == max(abs(G_new))*ones(size(G_new,1),1)) ...
%         - 1.02*min(rpinc,max(abs(G_new)))*(G_new < 0).*(abs(G_new) == max(abs(G_new))*ones(size(G_new,1),1));
%     supDifG = max(abs(G_new-G_old));
% end

% Put G_new on the grid 
rpinc = rp(2) - rp(1);
G_new_mod = mod(G_new,rpinc);
G_new_mod_up = (G_new_mod > rpinc/2);
G_new = G_new + G_new_mod_up.*(rpinc-G_new_mod) ...
    - (1-G_new_mod_up).*G_new_mod;

% Create distribution of S_t/P_t-1
%***************************************************

W_SPold = zeros(radgridsize,1);
for ii = 1:NSectors
    W_SPold = W_SPold + sum(reshape(Q(ii).Q,[agridsize*rpgridsize radgridsize]))';
end
W_SPold = (W_SPold >= 10^(-10)).*W_SPold;

